home *** CD-ROM | disk | FTP | other *** search
/ Personal Computer World 2008 February / PCWFEB08.iso / Software / Freeware / Miro 1.0 / Miro_Installer.exe / xulrunner / python / BitTorrent / PiecePicker.py < prev    next >
Encoding:
Python Source  |  2007-11-12  |  4.9 KB  |  182 lines
  1. # Written by Bram Cohen
  2. # see LICENSE.txt for license information
  3.  
  4. from random import randrange, shuffle, choice
  5.  
  6. class PiecePicker:
  7.     def __init__(self, numpieces, rarest_first_cutoff = 1):
  8.         self.rarest_first_cutoff = rarest_first_cutoff
  9.         self.numpieces = numpieces
  10.         self.interests = [range(numpieces)]
  11.         self.pos_in_interests = range(numpieces)
  12.         self.numinterests = [0] * numpieces
  13.         self.started = []
  14.         self.seedstarted = []
  15.         self.numgot = 0
  16.         self.scrambled = range(numpieces)
  17.         shuffle(self.scrambled)
  18.  
  19.     def got_have(self, piece):
  20.         if self.numinterests[piece] is None:
  21.             return
  22.         numint = self.numinterests[piece]
  23.         if numint == len(self.interests) - 1:
  24.             self.interests.append([])
  25.         self.numinterests[piece] += 1
  26.         self._shift_over(piece, self.interests[numint], self.interests[numint + 1])
  27.  
  28.     def lost_have(self, piece):
  29.         if self.numinterests[piece] is None:
  30.             return
  31.         numint = self.numinterests[piece]
  32.         self.numinterests[piece] -= 1
  33.         self._shift_over(piece, self.interests[numint], self.interests[numint - 1])
  34.  
  35.     def _shift_over(self, piece, l1, l2):
  36.         p = self.pos_in_interests[piece]
  37.         l1[p] = l1[-1]
  38.         self.pos_in_interests[l1[-1]] = p
  39.         del l1[-1]
  40.         newp = randrange(len(l2) + 1)
  41.         if newp == len(l2):
  42.             self.pos_in_interests[piece] = len(l2)
  43.             l2.append(piece)
  44.         else:
  45.             old = l2[newp]
  46.             self.pos_in_interests[old] = len(l2)
  47.             l2.append(old)
  48.             l2[newp] = piece
  49.             self.pos_in_interests[piece] = newp
  50.  
  51.     def requested(self, piece, seed = False):
  52.         if piece not in self.started:
  53.             self.started.append(piece)
  54.         if seed and piece not in self.seedstarted:
  55.             self.seedstarted.append(piece)
  56.  
  57.     def complete(self, piece):
  58.         assert self.numinterests[piece] is not None
  59.         self.numgot += 1
  60.         l = self.interests[self.numinterests[piece]]
  61.         p = self.pos_in_interests[piece]
  62.         l[p] = l[-1]
  63.         self.pos_in_interests[l[-1]] = p
  64.         del l[-1]
  65.         self.numinterests[piece] = None
  66.         try:
  67.             self.started.remove(piece)
  68.             self.seedstarted.remove(piece)
  69.         except ValueError:
  70.             pass
  71.  
  72.     def next(self, havefunc, seed = False):
  73.         bests = None
  74.         bestnum = 2 ** 30
  75.         if seed:
  76.             s = self.seedstarted
  77.         else:
  78.             s = self.started
  79.         for i in s:
  80.             if havefunc(i):
  81.                 if self.numinterests[i] < bestnum:
  82.                     bests = [i]
  83.                     bestnum = self.numinterests[i]
  84.                 elif self.numinterests[i] == bestnum:
  85.                     bests.append(i)
  86.         if bests:
  87.             return choice(bests)
  88.         if self.numgot < self.rarest_first_cutoff:
  89.             for i in self.scrambled:
  90.                 if havefunc(i):
  91.                     return i
  92.             return None
  93.         for i in xrange(1, min(bestnum, len(self.interests))):
  94.             for j in self.interests[i]:
  95.                 if havefunc(j):
  96.                     return j
  97.         return None
  98.  
  99.     def am_I_complete(self):
  100.         return self.numgot == self.numpieces
  101.  
  102.     def bump(self, piece):
  103.         l = self.interests[self.numinterests[piece]]
  104.         pos = self.pos_in_interests[piece]
  105.         del l[pos]
  106.         l.append(piece)
  107.         for i in range(pos,len(l)):
  108.             self.pos_in_interests[l[i]] = i
  109.  
  110. def test_requested():
  111.     p = PiecePicker(9)
  112.     p.complete(8)
  113.     p.got_have(0)
  114.     p.got_have(2)
  115.     p.got_have(4)
  116.     p.got_have(6)
  117.     p.requested(1)
  118.     p.requested(1)
  119.     p.requested(3)
  120.     p.requested(0)
  121.     p.requested(6)
  122.     v = _pull(p)
  123.     assert v[:2] == [1, 3] or v[:2] == [3, 1]
  124.     assert v[2:4] == [0, 6] or v[2:4] == [6, 0]
  125.     assert v[4:] == [2, 4] or v[4:] == [4, 2]
  126.  
  127. def test_change_interest():
  128.     p = PiecePicker(9)
  129.     p.complete(8)
  130.     p.got_have(0)
  131.     p.got_have(2)
  132.     p.got_have(4)
  133.     p.got_have(6)
  134.     p.lost_have(2)
  135.     p.lost_have(6)
  136.     v = _pull(p)
  137.     assert v == [0, 4] or v == [4, 0]
  138.  
  139. def test_change_interest2():
  140.     p = PiecePicker(9)
  141.     p.complete(8)
  142.     p.got_have(0)
  143.     p.got_have(2)
  144.     p.got_have(4)
  145.     p.got_have(6)
  146.     p.lost_have(2)
  147.     p.lost_have(6)
  148.     v = _pull(p)
  149.     assert v == [0, 4] or v == [4, 0]
  150.  
  151. def test_complete():
  152.     p = PiecePicker(1)
  153.     p.got_have(0)
  154.     p.complete(0)
  155.     assert _pull(p) == []
  156.     p.got_have(0)
  157.     p.lost_have(0)
  158.  
  159. def test_rarer_in_started_takes_priority():
  160.     p = PiecePicker(3)
  161.     p.complete(2)
  162.     p.requested(0)
  163.     p.requested(1)
  164.     p.got_have(1)
  165.     p.got_have(0)
  166.     p.got_have(0)
  167.     assert _pull(p) == [1, 0]
  168.  
  169. def test_zero():
  170.     assert _pull(PiecePicker(0)) == []
  171.  
  172. def _pull(pp):
  173.     r = []
  174.     def want(p, r = r):
  175.         return p not in r
  176.     while True:
  177.         n = pp.next(want)
  178.         if n is None:
  179.             break
  180.         r.append(n)
  181.     return r
  182.